package com.demo.function.dataauth.interceptor;

import com.demo.function.dataauth.manager.DataAuthManager;
import com.google.common.collect.Lists;
import lombok.extern.slf4j.Slf4j;
import net.sf.jsqlparser.JSQLParserException;
import net.sf.jsqlparser.expression.Expression;
import net.sf.jsqlparser.expression.operators.conditional.AndExpression;
import net.sf.jsqlparser.expression.operators.relational.EqualsTo;
import net.sf.jsqlparser.parser.CCJSqlParserUtil;
import net.sf.jsqlparser.schema.Table;
import net.sf.jsqlparser.statement.select.FromItem;
import net.sf.jsqlparser.statement.select.PlainSelect;
import net.sf.jsqlparser.statement.select.Select;
import net.sf.jsqlparser.statement.select.SelectBody;
import net.sf.jsqlparser.util.TablesNamesFinder;
import org.apache.commons.collections4.CollectionUtils;
import org.apache.commons.lang3.StringUtils;
import org.apache.ibatis.cache.CacheKey;
import org.apache.ibatis.executor.Executor;
import org.apache.ibatis.mapping.BoundSql;
import org.apache.ibatis.mapping.MappedStatement;
import org.apache.ibatis.mapping.ParameterMapping;
import org.apache.ibatis.plugin.Interceptor;
import org.apache.ibatis.plugin.Intercepts;
import org.apache.ibatis.plugin.Invocation;
import org.apache.ibatis.plugin.Signature;
import org.apache.ibatis.reflection.MetaObject;
import org.apache.ibatis.reflection.SystemMetaObject;
import org.apache.ibatis.session.Configuration;
import org.apache.ibatis.session.ResultHandler;
import org.apache.ibatis.session.RowBounds;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;

import javax.servlet.http.HttpServletRequest;
import java.util.List;
import java.util.Map;
import java.util.Objects;

@Slf4j
@Intercepts({
        @Signature(type = Executor.class, method = "query", args = {MappedStatement.class, Object.class, RowBounds.class, ResultHandler.class})
})
public class DataAuthInterceptor implements Interceptor {

    @Override
    public Object intercept(Invocation invocation) throws Throwable {
        log.info("数据权限拦截器");

        Executor executor = (Executor) invocation.getTarget();
        Object[] args = invocation.getArgs();
        MappedStatement ms = (MappedStatement) args[0];
        Object parameter = args[1];
        RowBounds rowBounds = (RowBounds) args[2];
        ResultHandler resultHandler = (ResultHandler) args[3];
        BoundSql boundSql = ms.getBoundSql(parameter);
        Select select = (Select) CCJSqlParserUtil.parse(boundSql.getSql());
        MetaObject boundSqlMetaObject = SystemMetaObject.forObject(boundSql);

        DataAuthManager instance = DataAuthManager.getInstance();
        ServletRequestAttributes attributes = (ServletRequestAttributes) RequestContextHolder.getRequestAttributes();
        if (attributes == null) {
            log.info("ServletRequestAttributes为空，不执行数据权限处理");
            return invocation.proceed();
        } else if (Objects.isNull(instance)) {
            log.info("DataAuthManager实例为空，不执行数据权限处理");
            return invocation.proceed();
        } else if (!isContainTable(select, instance)) {
            log.info("SQL不存在使用数据权限的表，不执行数据权限处理");
            return invocation.proceed();
        } else if (StringUtils.isBlank(attributes.getRequest().getHeader(instance.getHeaderKey()))) {
            log.info("请求头没有数据权限参数，不执行数据权限处理");
            return invocation.proceed();
        }

        setSql(boundSqlMetaObject, select, instance);
        addParameterMapping(boundSqlMetaObject, ms.getConfiguration(), instance);
        addAdditionalParameters(boundSql, attributes.getRequest(), instance);
        addParameterObject(parameter, attributes.getRequest(), instance);

        CacheKey cacheKey = executor.createCacheKey(ms, parameter, rowBounds, boundSql);
        return executor.query(ms, parameter, rowBounds, resultHandler, cacheKey, boundSql);
    }

    /**
     * sql语句是否包含需要数据权限的表
     * @param select jsql解析器select对象
     * @param instance DataAuthManager实例
     * @return true-数据权限的表存在
     */
    private boolean isContainTable(Select select, DataAuthManager instance) {
        TablesNamesFinder tablesNamesFinder = new TablesNamesFinder();
        List<String> tableList = tablesNamesFinder.getTableList(select);
        return tableList.contains(instance.getTableName());
    }

    /**
     * 重新生成sql
     * @param select jsql解析器select对象
     * @param instance DataAuthManager实例
     * @throws JSQLParserException 异常
     */
    private void setSql(MetaObject metaObject, Select select, DataAuthManager instance) throws JSQLParserException {
        SelectBody selectBody = select.getSelectBody();
        if (selectBody instanceof PlainSelect) {
            // 获取 From 模块
            PlainSelect plainSelect = (PlainSelect) selectBody;
            FromItem fromItem = plainSelect.getFromItem();
            if (fromItem instanceof Table) {
                Expression expression = plainSelect.getWhere();
                if (Objects.isNull(expression)) {
                    // 没有where条件，新增过滤条件
                    plainSelect.setWhere(new EqualsTo(CCJSqlParserUtil.parseCondExpression(instance.getColumnName()),
                            CCJSqlParserUtil.parseCondExpression("?")));
                } else {
                    // 在where条件后追加过滤条件
                    String filter = instance.getColumnName() + " = ?";
                    plainSelect.setWhere(new AndExpression(plainSelect.getWhere(), CCJSqlParserUtil.parseCondExpression(filter)));
                }
            }
        }
        String sql = select.toString();
        log.info("生成数据权限sql = " + sql);

        metaObject.setValue("dump.sql", sql);
    }

    /**
     * 添加参数名，这个名称对应参数注解@Param定义的value
     * @param metaObject BoundSql对象的MetaObject
     * @param configuration mybatis configuration对象
     * @param instance DataAuthManager实例
     */
    private void addParameterMapping(MetaObject metaObject, Configuration configuration,
                                     DataAuthManager instance) {
        String paramName = getParamName(instance);
        BoundSql boundSql = (BoundSql) metaObject.getOriginalObject();
        ParameterMapping parameterMapping = new ParameterMapping.Builder(configuration, paramName, Object.class).build();
        if (CollectionUtils.isEmpty(boundSql.getParameterMappings())) {
            metaObject.setValue("parameterMappings", Lists.newArrayList(parameterMapping));
        } else {
            boundSql.getParameterMappings().add(parameterMapping);
        }
    }

    private void addAdditionalParameters(BoundSql boundSql, HttpServletRequest request, DataAuthManager instance) {
        String paramName = getParamName(instance);
        String value = request.getHeader(instance.getHeaderKey());
        boundSql.setAdditionalParameter(paramName, value);
    }

    /**
     * 添加sql参数，key-boundSql.parameterMappings的值，value-替代sql占位符的值
     * @param parameter sql参数
     * @param request request对象
     * @param instance DataAuthManager实例
     */
    private void addParameterObject(Object parameter, HttpServletRequest request, DataAuthManager instance) {
        if (parameter == null) {
            return;
        }

        if (parameter instanceof Map) {
            String value = request.getHeader(instance.getHeaderKey());
            String paramName = getParamName(instance);

            Map parameterMap = (Map) parameter;
            parameterMap.put(paramName, value);
        }
    }

    private String getParamName(DataAuthManager instance) {
        return "dataAuth_" + instance.getColumnName();
    }
}
